#include "path_searching/dyn_a_star.h"

using namespace std;
using namespace Eigen;

AStar::~AStar()
{
}

void AStar::initGridMap(GridMap::Ptr occ_map, const Eigen::Vector3i pool_size)
{
  // init grid map and parameters
  grid_map_ = occ_map;
  grid_map_->getOrigin(orign_map_);
  reso_ = grid_map_->getResolution();
  inv_reso_ = 1.0 / reso_;

  POOL_SIZE_ = pool_size;
  CENTER_IDX_ = pool_size / 2;

  // init node map
  const size_t total_size = POOL_SIZE_(0) * POOL_SIZE_(1) * POOL_SIZE_(2);
  NodeMap_.reserve(total_size);
  for (size_t i = 0; i < total_size; i++)
    NodeMap_.emplace_back(std::make_shared<Node>());
}

vector<NodePtr> AStar::retrievePath(NodePtr current)
{
  vector<NodePtr> path;
  path.push_back(current);

  while (current->parent != nullptr)
  {
    current = current->parent;
    path.push_back(current);
  }

  return path;
}

/* -------------------------------------------------------------------------- */
/*            G cost and Heuristic Computation of H cost in A star            */
/* -------------------------------------------------------------------------- */

double AStar::getGCost(int dx, int dy, int dz)
{
  return sqrt(dx * dx + dy * dy + 16 * dz * dz);
}

/**
 * TODO MZY
 * 这个启发式计算H值和fast-planner有点不一样阿
 * 感觉因为这个是计算小范围内的路径，所以没必要那么复杂
 */
double AStar::getDiagHeu(NodePtr node1, NodePtr node2)
{
  double dx = abs(node1->index(0) - node2->index(0));
  double dy = abs(node1->index(1) - node2->index(1));
  double dz = abs(node1->index(2) - node2->index(2));

  double h = 0.0;
  int diag = min(min(dx, dy), dz);
  dx -= diag;
  dy -= diag;
  dz -= diag;

  if (dx == 0)
    h = 1.0 * sqrt(3.0) * diag + sqrt(2.0) * min(dy, dz) + 1.0 * abs(dy - dz);

  if (dy == 0)
    h = 1.0 * sqrt(3.0) * diag + sqrt(2.0) * min(dx, dz) + 1.0 * abs(dx - dz);

  if (dz == 0)
    h = 1.0 * sqrt(3.0) * diag + sqrt(2.0) * min(dx, dy) + 1.0 * abs(dx - dy);

  return h;
}

double AStar::getManhHeu(NodePtr node1, NodePtr node2)
{
  double dx = abs(node1->index(0) - node2->index(0));
  double dy = abs(node1->index(1) - node2->index(1));
  double dz = abs(node1->index(2) - node2->index(2));

  return dx + dy + dz;
}

double AStar::getEuclHeu(NodePtr node1, NodePtr node2)
{
  return (node2->index - node1->index).norm();
}

/**
 * @brief Check the legality of the input and perform preprocessing,
 *        update start_idx_,end_pt_
 */
bool AStar::checkInput(const Eigen::Vector3d &start_pt,
                       const Eigen::Vector3d &end_pt)
{
  if (start_pt == end_pt) {
    cerr << "\033[1;31m[Astar]: Start point is the same as end point, "
         << "force return! \033[0m" << endl;
    return false;
  }

  if (!setMapCenter((start_pt + end_pt) / 2.0)) {
    cerr << "\033[1;31m[Astar]: setMapCenter failed, "
         << "force return! \033[0m" << endl;
    return false;
  }

  if (!Pos2Index(start_pt, start_idx_) || !Pos2Index(end_pt, end_idx_)) {
    cout << "start_pt: " << start_pt.transpose() << endl;
    cout << "end_pt: " << end_pt.transpose() << endl;
    cerr << "\033[1;31m[Astar]: start/end pt is not in map, "
         << "force return! \033[0m" << endl;
    return false;
  }

  dir_ = end_pt - start_pt;
  dir_(2) = 0.0; // Dont change height of start and end point
  dir_ = dir_.normalized();
  step_ = dir_ * reso_;

  /* Extend on both ends to avoid endpoints within obstacles. */
  if (!expandPoint(start_pt, -step_, start_idx_) ||
      !expandPoint(end_pt, step_, end_idx_))
  {
    cerr << "\033[1;31m[Astar]: Unable to handle the endpoint"
         << " ,enlarge map size, force return! \033[0m" << endl;
    return false;
  }

  return true;
}

/* -------------------------------------------------------------------------- */
/*                                main function                               */
/* -------------------------------------------------------------------------- */

/**
 * @brief A-star path serach main function
 *        限定100*100*100 给定stepsize后，两点之间的最大距离就是100*stepsize，不然会出界
 */
int AStar::AstarSearch(const Eigen::Vector3d &start_pt,
                       const Eigen::Vector3d &end_pt)
{
  if (!checkInput(start_pt, end_pt)) return ERR_INVALIDINPUT;

  ++rounds_;

  // Initialize start and end node
  startPtr_ = NodeMap_[toAddress(start_idx_)];
  endPtr_ = NodeMap_[toAddress(end_idx_)];

  std::priority_queue<NodePtr, std::vector<NodePtr>, AstarComparator> empty;
  openSet_.swap(empty);

  startPtr_->index = start_idx_;
  startPtr_->rounds = rounds_;
  startPtr_->gScore = 0;
  startPtr_->fScore = getHCost(startPtr_, endPtr_);
  startPtr_->state = Node::OPENSET; // put start node in open set
  startPtr_->parent = nullptr;
  openSet_.push(startPtr_); // put start in open set

  endPtr_->index = end_idx_;

  // TODO 现在没有把search all的情况考虑进去
  // 如果是钻圈的话就需要考虑
  return search(false, false);
}

/**
 * @brief
 *
 * @param isGBFS
 * @param isSearchAll whether to search the entire map or
 *                    limit the height between start and end
 */
int AStar::search(bool isGBFS, bool isSearchAll)
{
  ros::Time time_begin = ros::Time::now();

  NodePtr neighborPtr = nullptr;
  NodePtr current = nullptr;

  int min_z = min(start_idx_(2), end_idx_(2));
  int max_z = max(start_idx_(2), end_idx_(2));

  // calculate the maximum number of iterations, x/y 2m more
  Eigen::Vector3i diff = (end_idx_ - start_idx_).cwiseAbs();
  int num_iter_max = int((diff(0) + 2 * inv_reso_) *
                         (diff(1) + 2 * inv_reso_) * diff(2));
  num_iter_max = std::max(num_iter_max, 1000);

  int num_iter = 0;
  while (!openSet_.empty())
  {
    num_iter++;
    current = openSet_.top();
    openSet_.pop();

    if (current->index(0) == endPtr_->index(0) &&
        current->index(1) == endPtr_->index(1) &&
        current->index(2) == endPtr_->index(2))
    {
      gridPath_ = retrievePath(current);

      if (gridPath_.size() == 1) return ERR_ONLYONEPOINT;

      return SUCCESS;
    }

    /* move current node from open set to closed set. */
    current->state = Node::CLOSEDSET;

    for (int dx = -1; dx <= 1; dx++)
      for (int dy = -1; dy <= 1; dy++)
        for (int dz = -1; dz <= 1; dz++)
        {
          if (dx == 0 && dy == 0 && dz == 0)
            continue;

          if (num_iter > num_iter_max) {
            // cerr << "\033[1;31m[Astar]: Failed in A star path searching !!!"
            //      << num_iter_max << " iter num limit exceeded. \033[0m" << endl;
            return ERR_MAXITERATION;
          }

          Vector3i neighborIdx;
          neighborIdx(0) = (current->index)(0) + dx;
          neighborIdx(1) = (current->index)(1) + dy;
          neighborIdx(2) = (current->index)(2) + dz;

          // Do not search the boundaries of the pool
          bool isOut;
          if (isSearchAll) {
            isOut = neighborIdx(0) < 1 || neighborIdx(0) >= POOL_SIZE_(0) - 1 ||
                    neighborIdx(1) < 1 || neighborIdx(1) >= POOL_SIZE_(1) - 1 ||
                    neighborIdx(2) < 1 || neighborIdx(2) >= POOL_SIZE_(2) - 1;
          }
          else {
            isOut = neighborIdx(0) < 1 || neighborIdx(0) >= POOL_SIZE_(0) - 1 ||
                    neighborIdx(1) < 1 || neighborIdx(1) >= POOL_SIZE_(1) - 1 ||
                    neighborIdx(2) < min_z || neighborIdx(2) > max_z;
          }
          if (isOut)
            continue;

          neighborPtr = NodeMap_[toAddress(neighborIdx)];
          neighborPtr->index = neighborIdx;

          bool flag_explored = neighborPtr->rounds == rounds_;

          if (flag_explored &&
              (neighborPtr->state == Node::CLOSEDSET ||
               neighborPtr->state == Node::OCCUPIED))
          {
            continue; // in closed set or occupied
          }
          else if (!flag_explored) {
            neighborPtr->state = Node::UNDEFINED;
            neighborPtr->rounds = rounds_;
          }

          int occ = checkOccupancy(Index2Pos(neighborPtr->index));
          if (occ == 1 || occ == -2) {
            neighborPtr->state = Node::OCCUPIED;
            continue; // occupied
          }

          // When isGBFS is true, gScore_tmp is not used
          double gScore_tmp;
          if (isGBFS)
            gScore_tmp = 0;
          else
            gScore_tmp = current->gScore + getGCost(dx, dy, dz);

          if (!flag_explored)
          { // discover a new node
            neighborPtr->state = Node::OPENSET;
            neighborPtr->parent = current;
            neighborPtr->gScore = gScore_tmp;
            neighborPtr->fScore = gScore_tmp + getHCost(neighborPtr, endPtr_);
            // put neighbor in open set and record it.
            openSet_.push(neighborPtr);
          }
          else if (gScore_tmp < neighborPtr->gScore)
          { // in open set and need update
            neighborPtr->parent = current;
            neighborPtr->gScore = gScore_tmp;
            neighborPtr->fScore = gScore_tmp + getHCost(neighborPtr, endPtr_);
          }
        }

    if ((ros::Time::now() - time_begin).toSec() > 0.2)
    {
      // cerr << "\033[1;31m[Astar]: Failed in A star path searching !!!"
      //      << "0.2 seconds time limit exceeded. \033[0m" << endl;
      return ERR_TIMEOUT;
    }
  }

  // cerr << "\033[1;31m[Astar]: Failed in A star path searching !!!"
  //      << "Open set is empty. \033[0m" << endl;
  return ERR_OPENSETEMPTY;
}

vector<Vector3d> AStar::getPath()
{
  size_t g_size = gridPath_.size();
  std::vector<Eigen::Vector3d> path(g_size);

  for (size_t i = 0; i < g_size; i++)
  {
    path[i] = Index2Pos(gridPath_[i]->index);
  }

  reverse(path.begin(), path.end());

  return path;
}
